from random import randint

import gym
import numpy as np
from gym.spaces import Discrete, Box


class EasyGameEnv(gym.Env):

    def __init__(self):
        self.action_space = Discrete(3)
        # 2x2 Box
        self.observation_space = Box(0, 1, (2, 2))
        self._max_episode_steps = 310

    def reset(self, *args, **kwargs):
        self.obs = np.array([[1, 0], [0, 0]])
        self.r = 0
        return self.obs

    def step(self, action):
        self.r += 1
        reward = 0
        done = False
        if action == 1:
            self.obs[0][0] = 1
            self.obs[0][1] = 0
            if self.obs[1][0] == 1:
                reward = -1
                done = True
        elif action == 2:
            reward = 0.001
            self.obs[0][1] = 1
            self.obs[0][0] = 0
            if self.obs[1][1] == 1:
                reward = -1
                done = True
        else:
            if self.obs[1][0] == self.obs[0][0] == 1:
                reward = -1
                done = True
            elif self.obs[1][1] == self.obs[0][1] == 1:
                reward = -1
                done = True
        if self.r > 300:
            done = True
            reward = 1
        i = randint(0, 1)
        next_one = [i, 1 - i]
        self.obs[1] = next_one
        return self.obs, reward, done, {}


if __name__ == '__main__':
    env = EasyGameEnv()
    env.reset()
    for i in range(1000):
        obs, reward, done, info = env.step(randint(0, 2))
        print(obs, reward, done, info)
        if done:
            env.reset()
            print('done')
